import multiprocessing as mp
import torch.distributed as dist
import torch.multiprocessing as tmp
from multiprocessing import shared_memory
import numpy as np
import time


class SharedMemoryManager:
    def __init__(self, world_size, rank):
        self.world_size = world_size  # 总进程数
        self.rank = rank  # 当前进程编号
        self.shm = None  # 共享内存对象
        self.data = None  # 共享内存中的数据

    def run(self):
        """主运行逻辑"""
        if self.world_size > 1:
            self.setup_shared_memory()

            if self.rank == 0:
                self.main_process_task()
            else:
                self.worker_process_task()

            self.cleanup()

    def setup_shared_memory(self):
        """设置共享内存和同步"""
        # 初始化分布式进程组
        dist.init_process_group(
            backend="gloo",
            init_method="tcp://localhost:23456",
            world_size=self.world_size,
            rank=self.rank
        )

        # 共享内存设置
        if self.rank == 0:
            # 主进程创建共享内存 (1MB)
            self.shm = shared_memory.SharedMemory(name="nanovllm", create=True, size=2 ** 20)
            print(f"主进程 {self.rank}: 创建共享内存 {self.shm.name}")
            dist.barrier()  # 等待所有进程准备就绪
        else:
            dist.barrier()  # 等待主进程创建共享内存
            # 子进程连接到已创建的共享内存
            self.shm = shared_memory.SharedMemory(name="nanovllm")
            print(f"子进程 {self.rank}: 连接到共享内存 {self.shm.name}")

        # 将共享内存映射为numpy数组以便操作 (假设存储100个整数)
        self.data = np.ndarray((250000,), dtype=np.int32, buffer=self.shm.buf)

    def main_process_task(self):
        """主进程任务：写入数据并通知子进程"""
        # 写入初始数据
        for i in range(10):
            self.data[i] = i * 100  # 写入示例数据
        print(f"主进程 {self.rank}: 写入初始数据: {self.data[:10]}")

        # 通知子进程可以读取数据
        self.data[250000 - 1] = 1  # 使用最后一个位置作为就绪标志
        time.sleep(1)

        # 等待子进程处理完成
        while self.data[250000 - 1] != 2:
            time.sleep(0.1)

        # 读取子进程处理后的数据
        print(f"主进程 {self.rank}: 读取子进程处理后的数据: {self.data[10:20]}")

    def worker_process_task(self):
        """子进程任务：读取数据、处理并返回结果"""
        # 等待主进程写入数据
        while self.data[250000 - 1] != 1:
            time.sleep(0.1)

        # 读取主进程写入的数据
        print(f"子进程 {self.rank}: 读取主进程数据: {self.data[:10]}")

        # 处理数据 (简单示例：加上进程编号)
        for i in range(10, 20):
            self.data[i] = self.data[i - 10] + self.rank

        # 通知主进程处理完成
        self.data[250000 - 1] = 2
        self.loop()  # 进入等待状态

    def loop(self):
        """子进程循环等待"""
        print(f"子进程 {self.rank}: 进入等待状态...")
        # 实际应用中可以等待新的任务信号
        time.sleep(2)

    def cleanup(self):
        """清理资源"""
        # 同步所有进程，确保数据处理完成
        dist.barrier()

        # 关闭共享内存
        self.shm.close()

        # 主进程负责销毁共享内存
        if self.rank == 0:
            self.shm.unlink()
            print(f"主进程 {self.rank}: 销毁共享内存")

        # 销毁进程组
        dist.destroy_process_group()


def run_process(rank, world_size):
    """进程入口函数"""
    manager = SharedMemoryManager(world_size, rank)
    manager.run()


if __name__ == "__main__":
    world_size = 3  # 总进程数：1个主进程 + 2个子进程

    # 启动多进程
    mp.set_start_method("spawn")
    processes = []
    for rank in range(world_size):
        p = mp.Process(target=run_process, args=(rank, world_size))
        p.start()
        processes.append(p)

    # 等待所有进程完成
    for p in processes:
        p.join()
